from torch import nn


class MultipleOutputLoss2(nn.Module):
    def __init__(self, weight_factors=None):
        """
        use this if you have several outputs and ground truth (both list of same len) and the loss should be computed
        between them (x[0] and y[0], x[1] and y[1] etc)
        :param loss:
        :param weight_factors:
        """
        super(MultipleOutputLoss2, self).__init__()
        self.weight_factors = weight_factors

    def cul_deep_loss(self, x):  # x 是loss集合的列表
        assert isinstance(x, (tuple, list)), "x must be either tuple or list"

        if self.weight_factors is None:
            weights = [1] * len(x)
        else:
            weights = self.weight_factors

        # l = weights[0] * self.loss(x[-1], y[0])
        l = 0
        for i in range(0, len(x)):
            if weights[i] != 0:
                l += weights[i] * x[i]
        return l
